{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "epochs = 15"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial 3 - Folded Split Neural Network\n",
    "\n",
    "\n",
    "<b>Recap:</b> The previous tutorial looked at building a SplitNN. The NN was split into three segments on three seperate hosts, where one host had data and another had labels. However, what if someone has data and labels in the same place? \n",
    "\n",
    "<b>Description: </b> Here we fold a multilayer SplitNN in on itself in order to accomodate the data nd labels being in the same place. We demonstrate the SplitNN class with a 3 segment distribution. This time,\n",
    "\n",
    "<img src=\"images/FoldedNN.png\" width=\"20%\">\n",
    "\n",
    "- <b>Alice</b>\n",
    "    - Has Model Segment 1\n",
    "    - Has Model Segment 3\n",
    "    - Has the handwritten images\n",
    "    - Has the image labels\n",
    "- <b>Bob</b>\n",
    "    - Has model Segment 2\n",
    "    \n",
    "Again, we use the exact same model as we used in the previous tutorial and see the same accuracy. Neither Alice nor Bob have the full model and Bob can't see Alice's data. \n",
    "\n",
    "Author:\n",
    "- Adam J Hall - Twitter: [@AJH4LL](https://twitter.com/AJH4LL) · GitHub:  [@H4LL](https://github.com/H4LL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SplitNN:\n",
    "    def __init__(self, models, optimizers):\n",
    "        self.models = models\n",
    "        self.optimizers = optimizers\n",
    "        \n",
    "    def forward(self, x):\n",
    "        a = []\n",
    "        remote_a = []\n",
    "        \n",
    "        a.append(models[0](x))\n",
    "        if a[-1].location == models[1].location:\n",
    "            remote_a.append(a[-1].detach().requires_grad_())\n",
    "        else:\n",
    "            remote_a.append(a[-1].detach().move(models[1].location).requires_grad_())\n",
    "\n",
    "        i=1    \n",
    "        while i < (len(models)-1):\n",
    "            \n",
    "            a.append(models[i](remote_a[-1]))\n",
    "            if a[-1].location == models[i+1].location:\n",
    "                remote_a.append(a[-1].detach().requires_grad_())\n",
    "            else:\n",
    "                remote_a.append(a[-1].detach().move(models[i+1].location).requires_grad_())\n",
    "            \n",
    "            i+=1\n",
    "        \n",
    "        a.append(models[i](remote_a[-1]))\n",
    "        self.a = a\n",
    "        self.remote_a = remote_a\n",
    "        \n",
    "        return a[-1]\n",
    "    \n",
    "    def backward(self):\n",
    "        a=self.a\n",
    "        remote_a=self.remote_a\n",
    "        optimizers = self.optimizers\n",
    "        \n",
    "        i= len(models)-2   \n",
    "        while i > -1:\n",
    "            if remote_a[i].location == a[i].location:\n",
    "                grad_a = remote_a[i].grad.copy()\n",
    "            else:\n",
    "                grad_a = remote_a[i].grad.copy().move(a[i].location)\n",
    "            a[i].backward(grad_a)\n",
    "            i-=1\n",
    "\n",
    "    \n",
    "    def zero_grads(self):\n",
    "        for opt in optimizers:\n",
    "            opt.zero_grad()\n",
    "        \n",
    "    def step(self):\n",
    "        for opt in optimizers:\n",
    "            opt.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torchvision import datasets, transforms\n",
    "from torch import nn, optim\n",
    "import syft as sy\n",
    "hook = sy.TorchHook(torch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Data preprocessing\n",
    "transform = transforms.Compose([transforms.ToTensor(),\n",
    "                              transforms.Normalize((0.5,), (0.5,)),\n",
    "                              ])\n",
    "trainset = datasets.MNIST('mnist', download=True, train=True, transform=transform)\n",
    "trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(0)\n",
    "\n",
    "# Define our model segments\n",
    "\n",
    "input_size = 784\n",
    "hidden_sizes = [128, 640]\n",
    "output_size = 10\n",
    "\n",
    "models = [\n",
    "    nn.Sequential(\n",
    "                nn.Linear(input_size, hidden_sizes[0]),\n",
    "                nn.ReLU(),\n",
    "    ),\n",
    "    nn.Sequential(\n",
    "                nn.Linear(hidden_sizes[0], hidden_sizes[1]),\n",
    "                nn.ReLU(),\n",
    "    ),\n",
    "    nn.Sequential(\n",
    "                nn.Linear(hidden_sizes[1], output_size),\n",
    "                nn.LogSoftmax(dim=1)\n",
    "    )\n",
    "]\n",
    "\n",
    "# Create optimisers for each segment and link to them\n",
    "optimizers = [\n",
    "    optim.SGD(model.parameters(), lr=0.03,)\n",
    "    for model in models\n",
    "]\n",
    "\n",
    "# create some workers\n",
    "alice = sy.VirtualWorker(hook, id=\"alice\")\n",
    "bob = sy.VirtualWorker(hook, id=\"bob\")\n",
    "\n",
    "# Send Model Segments to model locations\n",
    "model_locations = [alice, bob, alice]\n",
    "for model, location in zip(models, model_locations):\n",
    "    model.send(location)\n",
    "\n",
    "#Instantiate a SpliNN class with our distributed segments and their respective optimizers\n",
    "splitNN = SplitNN(models, optimizers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(x, target, splitNN):\n",
    "    \n",
    "    #1) Zero our grads\n",
    "    splitNN.zero_grads()\n",
    "    \n",
    "    #2) Make a prediction\n",
    "    pred = splitNN.forward(x)\n",
    "    \n",
    "    #3) Figure out how much we missed by\n",
    "    criterion = nn.NLLLoss()\n",
    "    loss = criterion(pred, target)\n",
    "    \n",
    "    #4) Backprop the loss on the end layer\n",
    "    loss.backward()\n",
    "    \n",
    "    #5) Feed Gradients backward through the nework\n",
    "    splitNN.backward()\n",
    "    \n",
    "    #6) Change the weights\n",
    "    splitNN.step()\n",
    "    \n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(epochs):\n",
    "    running_loss = 0\n",
    "    for images, labels in trainloader:\n",
    "        images = images.send(models[0].location)\n",
    "        images = images.view(images.shape[0], -1)\n",
    "        labels = labels.send(models[-1].location)\n",
    "        loss = train(images, labels, splitNN)\n",
    "        running_loss += loss.get()\n",
    "\n",
    "    else:\n",
    "        print(\"Epoch {} - Training loss: {}\".format(i, running_loss/len(trainloader)))"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
